#!/usr/bin/env python3

__author__ = 'Vladimir Iglovikov'

import math

import common.stats
import common.state

"""
Generates array of the states
nElectrons - number of electrons
sz - z component of the spin
nSites - number of sites
"""
def newStates(nElectrons, sz, nSites):
    nonFilteredNumberStates = int(math.pow(4, nSites))
    states = []
    for i in range(nonFilteredNumberStates):
        state = common.stats.int10to4array(i, nSites)
        tempState = common.state.State(state)
        if tempState.sz() == sz and tempState.nElectrons() == nElectrons:
            states += [tempState]
    return states

"""
For the set of states returns conversion table relating state with it's index in the Hamiltonian
states - array of the states we are worning with
"""
def conversionTable(states):
    nStates = len(states)
    result = []
    for i in range(nStates):
        result += [state2j(states[i])]
    result.sort()
    return result

"""
state - state that we are trying to represent
returns 2^nSites * I_down + I_up

I_down - binary representation of the downspins in a state
I_up - binary representation of the upspins in a state
"""
def state2j(state):
    nSites = state.nSites
    return int(math.pow(2, nSites) * state.binaryDown() + state.binaryUp())

"""
state - state that we are trying to represent
conTable - table relating binary representation wit hstate
returns index of the state in the Hamiltonian
"""
def state2I(state, conTable):
    result = conTable.search(state2j(state))
    return result